package com.leetcode.algorithm.y19.m04;

import org.junit.Assert;

import com.leetcode.algorithm.common.TreeNodeExt;

/**
 * 703. 数据流中的第K大元素
 * 
 * 设计一个找到数据流中第K大元素的类（class）。注意是排序后的第K大元素，不是第K个不同的元素。
 * 
 * 你的 KthLargest 类需要一个同时接收整数 k 和整数数组nums 的构造器，它包含数据流中的初始元素。每次调用 KthLargest.add，返回当前数据流中第K大的元素。
 * 
 * 示例:
 * 
 * int k = 3;
 * int[] arr = [4,5,8,2];
 * KthLargest kthLargest = new KthLargest(3, arr);
 * kthLargest.add(3);   // returns 4
 * kthLargest.add(5);   // returns 5
 * kthLargest.add(10);  // returns 5
 * kthLargest.add(9);   // returns 8
 * kthLargest.add(4);   // returns 8
 * 说明: 
 * 你可以假设 nums 的长度≥ k-1 且k ≥ 1。
 * @author: jie.deng
 * @time: 2019年4月8日 下午7:28:24
 */
class KthLargest {

    private TreeNodeExt root;
    private int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        for (int num : nums) {
            addNode(num);
        }
    }

    public int add(int val) {
        addNode(val);
        return searchKthLargest(this.k);
    }

    public int searchKthLargest(int k) {
        TreeNodeExt pre = this.root;
        TreeNodeExt cur = null;
        while (k > 0) {
            cur = pre.right;
            while (cur != null && cur.i1 >= k) {
                pre = cur;
                cur = pre.right;
            }
            if (cur != null) {
                k -= cur.i1; // 右子树
            }
            k -= 1; // 当前结点
            if (k >= 1) {
                pre = pre.left; // 左子树
            }
        }
        return pre.val;
    }

    /**
     * 插入结点
     * 
     * @param val
     */
    private void addNode(int val) {
        TreeNodeExt node = new TreeNodeExt(val, 1);
        if (this.root == null) {
            this.root = node;
            return;
        }
        TreeNodeExt pre = null;
        TreeNodeExt cur = this.root;
        while (cur != null && cur.val != val) {
            pre = cur;
            pre.i1++;
            if (cur.val > val) {
                cur = cur.left;
            } else if (cur.val < val) {
                cur = cur.right;
            }
        }

        if (cur == null) {
            // 新增结点是叶子结点，作为pre的孩子结点
            if (pre == null) {
                this.root = node;
            } else {
                if (pre.val > val) {
                    pre.left = node;
                } else {
                    pre.right = node;
                }
            }
        } else {
            // 如果与当前结点key值相同，作为当前结点的右孩子
            node.right = cur.right;
            node.i1 += cur.right == null ? 0 : cur.right.i1;
            cur.right = node;
            cur.i1++;
        }
    }
    
    public static void main(String[] args) {
        int k = 3;
        int[] arr = new int[] { 4, 5, 8, 2 };
        KthLargest kthLargest = new KthLargest(k, arr);
        Assert.assertEquals(kthLargest.add(3), 4);// returns 4
        Assert.assertEquals(kthLargest.add(5), 5); // returns 5
        Assert.assertEquals(kthLargest.add(10), 5); // returns 5
        Assert.assertEquals(kthLargest.add(9), 8); // returns 8
        Assert.assertEquals(kthLargest.add(4), 8); // returns 8
        
        k = 7;
        arr = new int[] { -10, 1, 3, 1, 4, 10, 3, 9, 4, 5, 1 };
        kthLargest = new KthLargest(k, arr);
        Assert.assertEquals(kthLargest.add(3), 3);
        Assert.assertEquals(kthLargest.add(2), 3);
        Assert.assertEquals(kthLargest.add(3), 3);
        Assert.assertEquals(kthLargest.add(1), 3);
        Assert.assertEquals(kthLargest.add(2), 3);
        Assert.assertEquals(kthLargest.add(4), 3);
        Assert.assertEquals(kthLargest.add(5), 4);
        Assert.assertEquals(kthLargest.add(5), 4);
        Assert.assertEquals(kthLargest.add(6), 4);
        Assert.assertEquals(kthLargest.add(7), 5);
        Assert.assertEquals(kthLargest.add(7), 5);
        Assert.assertEquals(kthLargest.add(8), 5);
        Assert.assertEquals(kthLargest.add(2), 5);
        Assert.assertEquals(kthLargest.add(3), 5);
        Assert.assertEquals(kthLargest.add(1), 5);
        Assert.assertEquals(kthLargest.add(1), 5);
        Assert.assertEquals(kthLargest.add(1), 5);
        Assert.assertEquals(kthLargest.add(10), 6);
        Assert.assertEquals(kthLargest.add(11), 7);
        Assert.assertEquals(kthLargest.add(5), 7);
        Assert.assertEquals(kthLargest.add(6), 7);
        Assert.assertEquals(kthLargest.add(2), 7);
        Assert.assertEquals(kthLargest.add(4), 7);
        Assert.assertEquals(kthLargest.add(7), 7);
        Assert.assertEquals(kthLargest.add(8), 7);
        Assert.assertEquals(kthLargest.add(5), 7);
        Assert.assertEquals(kthLargest.add(6), 7);
    }
}

/**
 * Your KthLargest object will be instantiated and called as such:
 * KthLargest obj = new KthLargest(k, nums);
 * int param_1 = obj.add(val);
 */